#include "KdTree.h"
#include "STLInclude.h"
#include <tchar.h>
#include <assert.h>

bool RayIntersectTriangle(const RAY& ray, const Triangle& triangle, float& u, float& v, float& distance)
{
	bool hitResult = false;
	VEC3 intersectionPoint;
	float d = ray.direction.Dot(triangle.plane.Normal);

	if (d > CRT_MATH_EPSILON || d < -CRT_MATH_EPSILON)
	{
		distance = (-triangle.plane.Distance - ray.origin.Dot(triangle.plane.Normal)) / d;
		
		if (distance >= 0)
		{
			intersectionPoint = ray.origin + ray.direction * distance;

			u = intersectionPoint.Dot(triangle.vecPerpToV02) + triangle.b1;

			if (u >= 0 && u <= 1.0f)
			{
				v = intersectionPoint.Dot(triangle.vecPerpToV01) + triangle.b2;

				hitResult = (v >= 0) && (v <= 1.0f) && (u + v <= 1.0f);
			} 
		}	
	}

	return hitResult;
}

// If a point is "on" the plane, it's considered to be at the front side of the plane
static inline bool PointAtFrontSideOfPlane(const VEC3& point, int axis, float split)
{
	return ((float*)&point)[axis] <= split;
}

// Kay/Kajiya slabs test
// http://tog.acm.org/resources/RTNews/html/rtnv19n1.html#art7
bool RayIntersectAABB(const RAY& ray, const AABB& aabb)
{
    float tmin, tmax, tymin, tymax, tzmin, tzmax;

    tmin = (((VEC3*)&aabb)[ray.sign[0]].x - ray.origin.x) * ray.invDirection.x;
    tmax = (((VEC3*)&aabb)[1-ray.sign[0]].x - ray.origin.x) * ray.invDirection.x;
    tymin = (((VEC3*)&aabb)[ray.sign[1]].y - ray.origin.y) * ray.invDirection.y;
    tymax = (((VEC3*)&aabb)[1-ray.sign[1]].y - ray.origin.y) * ray.invDirection.y;
    if ( (tmin > tymax) || (tymin > tmax) ) // branch 1 and branch 2
        return false;
    if (tymin > tmin)
        tmin = tymin;
    if (tymax < tmax)
        tmax = tymax;
    tzmin = (((VEC3*)&aabb)[ray.sign[2]].z - ray.origin.z) * ray.invDirection.z;
    tzmax = (((VEC3*)&aabb)[1-ray.sign[2]].z - ray.origin.z) * ray.invDirection.z;
    if ( (tmin > tzmax) || (tzmin > tmax) ) // branch 3 and branch 4
        return false;
	return true;
    //if (tzmin > tmin)
    //    tmin = tzmin;
    //if (tzmax < tmax)
    //    tmax = tzmax;
    //// Note: I have changed the end condition here to fit standard practice
    //return ( (tmin <= t1) && (tmax >= t0) ); // branch 5 (fail tmin), 6 (fail tmax), and 7 (success)
}

KdTreeNode::KdTreeNode()
{
	mpFrontChild = NULL;
	mpBackChild = NULL;
	mIsLeaf = true;
	mAxis = 0;
	mSplit = 0;
	mTriangleList.clear();
}

KdTreeNode::~KdTreeNode()
{
	CleanMembers();
}

void KdTreeNode::CleanMembers()
{
	delete mpFrontChild;
	mpFrontChild = NULL;

	delete mpBackChild;
	mpBackChild = NULL;

	mTriangleList.clear();
}

void CompactKdTreeNode::TraceRay(const KdTree* pTree, RAY* pRay, TraceType traceType, ITraceResult* pResult) const
{
	// First of all, test ray against aabb
	bool hitAABB = RayIntersectAABB(const_cast<const RAY&>(*pRay), mAABB);
	if (NULL != pResult)
	{
		pResult->AddRayAABBTesting();
	}

	if (!hitAABB)
	{
		return;
	}

	if (IsLeaf())
	{
		unsigned int* pTriangleList = TriangleList(pTree);
		for (size_t i = 0; i < TriangleNum(); i++)
		{
			int triangleIndex = pTriangleList[i];
			if (pTree->mpTriangleVisitedFlag[triangleIndex])
			{
				continue;
			}

			Triangle& triangle = pTree->mpTriangleList[triangleIndex];
			float u, v, distance;

			bool bHit = RayIntersectTriangle(*pRay, triangle, u, v, distance);
			pTree->mpTriangleVisitedFlag[triangleIndex] = true;

			if (NULL != pResult)
			{
				pResult->AddRayTriangleTesting();

				if (bHit)
				{
					TSceneIntersection intersection;
					intersection.u = u;
					intersection.v = v;
					intersection.dist = distance;
					intersection.meshIndex = triangle.meshIndex;

					intersection.index[0] = triangle.vertexIndex[0] - pTree->mpVertexOffsetList[triangleIndex];
					intersection.index[1] = triangle.vertexIndex[1] - pTree->mpVertexOffsetList[triangleIndex];
					intersection.index[2] = triangle.vertexIndex[2] - pTree->mpVertexOffsetList[triangleIndex];

					pResult->AddHit(intersection, traceType);

					if (pResult->CanStopTrace())
					{
						return;
					}
				}
			}
		}
	}
	else
	{
		int iAxis = Axis();
		float fSplit = Split();
		CompactKdTreeNode* pFrontChild = FrontChild(pTree);
		CompactKdTreeNode* pBackChild = BackChild(pTree);

		bool rayOriginAtFrontSideOfSplittingPlane = 
			PointAtFrontSideOfPlane(pRay->effectiveOrigin, iAxis, fSplit);
		bool rayTargetAtFrontSideOfSplittingPlane = 
			PointAtFrontSideOfPlane(pRay->effectiveTarget, iAxis, fSplit);

		if (rayOriginAtFrontSideOfSplittingPlane && rayTargetAtFrontSideOfSplittingPlane)
		{
			pFrontChild->TraceRay(pTree, pRay, traceType, pResult);
		}
		else if (!rayOriginAtFrontSideOfSplittingPlane && !rayTargetAtFrontSideOfSplittingPlane)
		{
			pBackChild->TraceRay(pTree, pRay, traceType, pResult);
		}
		else
		{
			// Find the intersection point
			VEC3 intersectionPoint;

			float fStart = ((float*)&pRay->effectiveOrigin)[iAxis];
			float fEnd = ((float*)&pRay->effectiveTarget)[iAxis];

			float ratio = (fSplit - fStart) / (fEnd - fStart);
			intersectionPoint = pRay->effectiveOrigin * (1.0f - ratio) + pRay->effectiveTarget * ratio;

			VEC3 oldTarget = pRay->effectiveTarget;
			pRay->effectiveTarget = intersectionPoint;

			if (rayOriginAtFrontSideOfSplittingPlane)
			{
				pFrontChild->TraceRay(pTree, pRay, traceType, pResult);
			}
			else
			{
				pBackChild->TraceRay(pTree, pRay, traceType, pResult);
			}

			if (kStopOnFirstHit == traceType && pResult->HitAny())
			{
				return;
			}
			
			if (kNearestHit == traceType && pResult->HitAny())
			{
				const TSceneIntersection* pIntersection = pResult->GetHit(0);
				VEC3 intersectionP = pRay->origin + pRay->direction * pIntersection->dist;
				bool intersectionAtFrontSideOfSplittingPlane = 
					PointAtFrontSideOfPlane(intersectionP, iAxis, fSplit);

				if (intersectionAtFrontSideOfSplittingPlane == rayOriginAtFrontSideOfSplittingPlane)
				{
					return;
				}
			}

			pRay->effectiveOrigin = intersectionPoint;
			pRay->effectiveTarget = oldTarget;

			if (rayTargetAtFrontSideOfSplittingPlane)
			{
				pFrontChild->TraceRay(pTree, pRay, traceType, pResult);
			}
			else
			{
				pBackChild->TraceRay(pTree, pRay, traceType, pResult);
			}
		}
	}
}


void KdTreeNode::TraceRay(const KdTree* pTree, RAY* pRay, TraceType traceType, ITraceResult* pResult) const
{
	// First of all, test ray against aabb
	bool hitAABB = RayIntersectAABB(const_cast<const RAY&>(*pRay), mAABB);
	if (NULL != pResult)
	{
		pResult->AddRayAABBTesting();
	}

	if (!hitAABB)
	{
		return;
	}

	if (mIsLeaf)
	{
		for (size_t i = 0; i < mTriangleList.size(); i++)
		{
			int triangleIndex = mTriangleList[i];
			if (pTree->mpTriangleVisitedFlag[triangleIndex])
			{
				continue;
			}

			Triangle& triangle = pTree->mpTriangleList[triangleIndex];
			float u, v, distance;

			bool bHit = RayIntersectTriangle(*pRay, triangle, u, v, distance);
			pTree->mpTriangleVisitedFlag[triangleIndex] = true;

			if (NULL != pResult)
			{
				pResult->AddRayTriangleTesting();

				if (bHit)
				{
					TSceneIntersection intersection;
					intersection.u = u;
					intersection.v = v;
					intersection.dist = distance;
					intersection.meshIndex = triangle.meshIndex;

					intersection.index[0] = triangle.vertexIndex[0] - pTree->mpVertexOffsetList[mTriangleList[i]];
					intersection.index[1] = triangle.vertexIndex[1] - pTree->mpVertexOffsetList[mTriangleList[i]];
					intersection.index[2] = triangle.vertexIndex[2] - pTree->mpVertexOffsetList[mTriangleList[i]];

					pResult->AddHit(intersection, traceType);

					if (pResult->CanStopTrace())
					{
						return;
					}
				}
			}
		}
	}
	else
	{
		bool rayOriginAtFrontSideOfSplittingPlane = 
			PointAtFrontSideOfPlane(pRay->effectiveOrigin, mAxis, mSplit);
		bool rayTargetAtFrontSideOfSplittingPlane = 
			PointAtFrontSideOfPlane(pRay->effectiveTarget, mAxis, mSplit);

		if (rayOriginAtFrontSideOfSplittingPlane && rayTargetAtFrontSideOfSplittingPlane)
		{
			mpFrontChild->TraceRay(pTree, pRay, traceType, pResult);
		}
		else if (!rayOriginAtFrontSideOfSplittingPlane && !rayTargetAtFrontSideOfSplittingPlane)
		{
			mpBackChild->TraceRay(pTree, pRay, traceType, pResult);
		}
		else
		{
			// Find the intersection point
			VEC3 intersectionPoint;

			float fStart = ((float*)&pRay->effectiveOrigin)[mAxis];
			float fEnd = ((float*)&pRay->effectiveTarget)[mAxis];

			float ratio = (mSplit - fStart) / (fEnd - fStart);
			intersectionPoint = pRay->effectiveOrigin * (1.0f - ratio) + pRay->effectiveTarget * ratio;

			VEC3 oldTarget = pRay->effectiveTarget;
			pRay->effectiveTarget = intersectionPoint;

			if (rayOriginAtFrontSideOfSplittingPlane)
			{
				mpFrontChild->TraceRay(pTree, pRay, traceType, pResult);
			}
			else
			{
				mpBackChild->TraceRay(pTree, pRay, traceType, pResult);
			}

			if (kStopOnFirstHit == traceType && pResult->HitAny())
			{
				return;
			}
			
			if (kNearestHit == traceType && pResult->HitAny())
			{
				const TSceneIntersection* pIntersection = pResult->GetHit(0);
				VEC3 intersectionP = pRay->origin + pRay->direction * pIntersection->dist;
				bool intersectionAtFrontSideOfSplittingPlane = 
					PointAtFrontSideOfPlane(intersectionP, mAxis, mSplit);

				if (intersectionAtFrontSideOfSplittingPlane == rayOriginAtFrontSideOfSplittingPlane)
				{
					return;
				}
			}

			pRay->effectiveOrigin = intersectionPoint;
			pRay->effectiveTarget = oldTarget;

			if (rayTargetAtFrontSideOfSplittingPlane)
			{
				mpFrontChild->TraceRay(pTree, pRay, traceType, pResult);
			}
			else
			{
				mpBackChild->TraceRay(pTree, pRay, traceType, pResult);
			}
		}
	}
}

void KdTreeNode::GetLocalBoundingBox(const KdTree* pTree, AABB& bb)
{
	bb.Init();

	for (size_t i = 0; i < mTriangleList.size(); i++)
	{
		Triangle& triangle = pTree->mpTriangleList[mTriangleList[i]];
		
		bb += pTree->mpVertexList[triangle.vertexIndex[0]].position;
		bb += pTree->mpVertexList[triangle.vertexIndex[1]].position;
		bb += pTree->mpVertexList[triangle.vertexIndex[2]].position;
	}
}

size_t KdTreeNode::GetTriangleCount() const
{
	if (mIsLeaf)
	{
		return mTriangleList.size();
	}
	else
	{
		return mpFrontChild->GetTriangleCount() + 
			mpBackChild->GetTriangleCount();
	}
}



size_t KdTreeNode::GetNumberOfNode() const
{
	if (mIsLeaf)
	{
		return 1;
	}
	else
	{
		return 1 + mpFrontChild->GetNumberOfNode() + mpBackChild->GetNumberOfNode();
	}
}


void KdTreeNode::BuildCompact(
		KdTree* pTree, 
		CompactKdTreeNode* pCurrentNode,
		unsigned int& iAvailableNode,
		unsigned int& iAvailableLeafData
		)
{
	pCurrentNode->mMultiUsageFlag = 0;
	pCurrentNode->mMultiUsageFlag |= mIsLeaf ? 0x80000000 : 0;
	pCurrentNode->mAABB = mAABB;

	if (mIsLeaf)
	{
		pCurrentNode->mMultiUsageFlag |= iAvailableLeafData & 0x1fffffff;

		unsigned int* pLeafTriangleData = pTree->mpLeafTriangleData + iAvailableLeafData;
		pCurrentNode->mU.mTriangleNum = (unsigned int)mTriangleList.size();
		for (unsigned int i = 0; i < pCurrentNode->mU.mTriangleNum; i++)
		{
			pLeafTriangleData[i] = mTriangleList[i];
		}

		iAvailableLeafData += mTriangleList.size();
	}
	else
	{
		pCurrentNode->mMultiUsageFlag |= (mAxis << 29) & 0x70000000;
		pCurrentNode->mMultiUsageFlag |= iAvailableNode & 0x1fffffff;
		pCurrentNode->mU.mSplit = mSplit;

		CompactKdTreeNode* pFrontChild = pTree->mpCompactRoot + iAvailableNode;
		CompactKdTreeNode* pBackChild = pFrontChild + 1;
		iAvailableNode += 2;

		mpFrontChild->BuildCompact(pTree, pFrontChild, iAvailableNode, iAvailableLeafData);
		mpBackChild->BuildCompact(pTree, pBackChild, iAvailableNode, iAvailableLeafData);
	}
}



KdTree::KdTree(const vector<TCRTMesh*>& Meshes)
{
	mpRoot = NULL;
	mpVertexList = NULL;
	mpTriangleList = NULL;
	mpVertexOffsetList = NULL;
	mpTriangleVisitedFlag = NULL;

	mVertexCount = 0;
	mTriangleCount = 0;

	mRefMeshes.empty();
	for (size_t i = 0; i < Meshes.size(); i++)
	{
		mRefMeshes.push_back(Meshes[i]);
	}
	mTreeReady = false;

	mpCompactRoot = NULL;
	mpLeafTriangleData = NULL;
	mNodeNum = 0;
	mLeafTriangleNum = 0;
}

KdTree::~KdTree()
{
	CleanMembers();
	mRefMeshes.empty();
}

bool KdTree::IsReady() const
{
	return mTreeReady;
}

void KdTree::CleanMembers()
{
	delete []mpVertexList;
	mpVertexList = NULL;

	delete []mpTriangleList;
	mpTriangleList = NULL;

	delete []mpVertexOffsetList;
	mpVertexOffsetList = NULL;

	delete []mpTriangleVisitedFlag;
	mpTriangleVisitedFlag = NULL;

	mVertexCount = 0;
	mTriangleCount = 0;

	delete mpRoot;
	mpRoot = NULL;

	mTreeReady = false;

	delete []mpCompactRoot;
	mpCompactRoot = NULL;

	delete []mpLeafTriangleData;
	mpLeafTriangleData = NULL;

	mLeafTriangleNum = 0;
	mNodeNum = 0;
}

bool KdTree::Build()
{
	if (mTreeReady)
	{
		return true;
	}

	CleanMembers();

	// Data copy
	mTriangleCount = 0;
	mVertexCount = 0;
	size_t nMesh = mRefMeshes.size();

	for (size_t i = 0; i < nMesh; i++)
	{
		const TCRTMesh* pMesh = mRefMeshes[i];
		mVertexCount += pMesh->vertices_f.size();
		mTriangleCount += pMesh->triangles.size();
	}

	mpVertexList = new Vertex[mVertexCount];
	mpTriangleList = new Triangle[mTriangleCount];
	mpVertexOffsetList = new int[mTriangleCount];
	mpTriangleVisitedFlag = new bool[mTriangleCount];
	memset(mpTriangleVisitedFlag, 0, sizeof(bool) * mTriangleCount);

	Vertex* pVertex = mpVertexList;
	Triangle* pTriangle = mpTriangleList;
	int* pVertexOffset = mpVertexOffsetList;

	size_t nVertexOffset = 0;
	size_t nTriangle = 0;
	size_t nVertex = 0;
	for (size_t i = 0; i < nMesh; i++)
	{
		const TCRTMesh* pMesh = mRefMeshes[i];
		nTriangle = pMesh->triangles.size();
		nVertex = pMesh->vertices_f.size();

		for (size_t j = 0; j < nVertex; j++)
		{
			(*pVertex++).position = pMesh->vertices_f[j].position;
		}
		
		for (size_t j = 0; j < nTriangle; j++)
		{
			const TCRTTriangle& crtTriangle = pMesh->triangles[j];
			const TCRTTriangleHelper& crtTriangleHelper = pMesh->triangleHelpers[j];

			pTriangle->meshIndex = (int)i;
			pTriangle->vertexIndex[0] = crtTriangle.index[0] + nVertexOffset;
			pTriangle->vertexIndex[1] = crtTriangle.index[1] + nVertexOffset;
			pTriangle->vertexIndex[2] = crtTriangle.index[2] + nVertexOffset;
			pTriangle->plane = crtTriangleHelper.plane;
			pTriangle->vecPerpToV01 = crtTriangleHelper.perp1;
			pTriangle->vecPerpToV02 = crtTriangleHelper.perp2;
			pTriangle->b1 = pMesh->triangleHelpers[j].b1;
			pTriangle->b2 = pMesh->triangleHelpers[j].b2;
			pTriangle++;

			(*pVertexOffset++) = nVertexOffset;
		}

		nVertexOffset += nVertex;
	}

	// Create root node
	mpRoot = new KdTreeNode();
	mpRoot->mIsLeaf = true;
	mpRoot->mTriangleList.resize(mTriangleCount);
	for (size_t i = 0; i < mTriangleCount; i++)
	{
		mpRoot->mTriangleList[i] = i;
	}

	// Build tree
	mMaximumDepth = 0;
	AABB bb;
	mpRoot->GetLocalBoundingBox(this, bb);
	BuildSubTree(mpRoot, bb, 0);

	// Build compact tree
	mNodeNum = mpRoot->GetNumberOfNode();
	mLeafTriangleNum = mpRoot->GetTriangleCount();

	mpCompactRoot = new CompactKdTreeNode[mNodeNum];
	mpLeafTriangleData = new unsigned int[mLeafTriangleNum];

	unsigned int iAvailableNode = 1;
	unsigned int iAvailableLeafData = 0;
	mpRoot->BuildCompact(this, mpCompactRoot, iAvailableNode, iAvailableLeafData);
	assert(iAvailableNode == mNodeNum);
	assert(iAvailableLeafData == mLeafTriangleNum);
	ReportStats();

	mTreeReady = true;
	return true;
}

void KdTree::BuildSubTree(KdTreeNode* pRoot, AABB& bb, size_t iLevel)
{
	if (NULL == pRoot)
	{
		return;
	}

	pRoot->mAABB = bb;
	
	if (StopSplitting(pRoot, iLevel))
	{
		return;
	}

	if (iLevel > mMaximumDepth)
	{
		mMaximumDepth = iLevel;
	}

	int splitAxis;
	float splitPos;

	if (FindSplitPlane(pRoot, iLevel, splitAxis, splitPos))
	{
		pRoot->mIsLeaf = false;
		pRoot->mpFrontChild = new KdTreeNode();
		pRoot->mpBackChild = new KdTreeNode();
		pRoot->mAxis = splitAxis;
		pRoot->mSplit = splitPos;

		for (size_t i = 0; i < pRoot->mTriangleList.size(); i++)
		{
			int triangleIndex = pRoot->mTriangleList[i];
			Triangle& triangle = mpTriangleList[triangleIndex];
			Vertex& v0 = mpVertexList[triangle.vertexIndex[0]];
			Vertex& v1 = mpVertexList[triangle.vertexIndex[1]];
			Vertex& v2 = mpVertexList[triangle.vertexIndex[2]];

			bool frontSide = false;
			bool backSide = false;

			if (PointAtFrontSideOfPlane(v0.position, splitAxis, splitPos))
			{
				frontSide = true;
			}
			else 
			{
				backSide = true;
			}

			if (PointAtFrontSideOfPlane(v1.position, splitAxis, splitPos))
			{
				frontSide = true;
			}
			else 
			{
				backSide = true;
			}

			if (PointAtFrontSideOfPlane(v2.position, splitAxis, splitPos))
			{
				frontSide = true;
			}
			else 
			{
				backSide = true;
			}

			if (frontSide)
			{
				pRoot->mpFrontChild->mTriangleList.push_back(triangleIndex);
			}
			if (backSide)
			{
				pRoot->mpBackChild->mTriangleList.push_back(triangleIndex);
			}
		}

		pRoot->mTriangleList.clear();

		AABB halfBB = bb;
		AABB childBB;

		((float*)&halfBB.maxV)[splitAxis] = splitPos;
		pRoot->mpFrontChild->GetLocalBoundingBox(this, childBB);
		childBB.Intersect(halfBB, childBB);
		
		BuildSubTree(pRoot->mpFrontChild, childBB, iLevel + 1);

		halfBB = bb;
		((float*)&halfBB.minV)[splitAxis] = splitPos;
		pRoot->mpBackChild->GetLocalBoundingBox(this, childBB);
		childBB.Intersect(halfBB, childBB);

		BuildSubTree(pRoot->mpBackChild, childBB, iLevel + 1);
	}
}

void KdTree::TraceRay(RAY* pRay, TraceType traceType, ITraceResult* pResult)
{
	if (NULL == mpRoot)
	{
		return;
	}

	// Clear the flag buffer
	memset(mpTriangleVisitedFlag, 0, sizeof(bool) * mTriangleCount);

	mpCompactRoot->TraceRay(this, pRay, traceType, pResult);
	//mpRoot->TraceRay(this, pRay, traceType, pResult);
}

void KdTree::ReportStats() const
{
	TCHAR stat[1024];

	_stprintf_s(stat, 1024, _T("Maximum depth: %d\n"), mMaximumDepth);
	OutputDebugString(stat);

	size_t triangleCount = mpRoot->GetTriangleCount();
	_stprintf_s(stat, 1024, _T("Tree efficiency: %.2f\n"), 1.0f * triangleCount / mTriangleCount);
	OutputDebugString(stat);
}